package common.test;

import common.TreeNode;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * 236. 二叉树的最近公共祖先
 */
public class _236_LowestCommonAncestor {

    Map<Integer, TreeNode> parent = new HashMap<>();
    Set<TreeNode> set = new HashSet<>();

    public void dfs(TreeNode root) {
        if (root == null) {
            return;
        }
        if (root.left != null) {
            parent.put(root.left.val, root);
            dfs(root.left);
        }
        if (root.right != null) {
            parent.put(root.right.val, root);
            dfs(root.right);
        }
    }


    public TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
        dfs(root);
        while (p != null) {
            set.add(p);
            p = parent.get(p.val);
        }
        while (q != null) {
            if (set.contains(q)) {
                return q;
            }
            q = parent.get(q.val);
        }
        return null;
    }

    TreeNode ans = null;

    public TreeNode lowestCommonAncestor2(TreeNode root, TreeNode p, TreeNode q) {
        if (root == null) {
            return null;
        }
        dfs(root, p, q);
        return ans;
    }


    public boolean dfs(TreeNode root, TreeNode p, TreeNode q) {
        if (root == null) {
            return false;
        }
        boolean l = dfs(root.left, p, q);
        boolean r = dfs(root.right, p, q);
        if ((l && r) || ((root.val == p.val || root.val == q.val) && (l || r))) {
            ans = root;
        }
        return l || r || (root.val == p.val || root.val == q.val);
    }
}
